import os
import numpy as np
import argparse
import sys  
sys.path.append(sys.path[0] + '/../caffe-onnx')
from proto import caffe_nvt_pb2 as caffe_pb2

def process_command():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config_dir', '-c', type=str, required=True, help='input config path')
    parser.add_argument('--input', '-i', type=str, required=False, help='model input path')
    parser.add_argument('--name', '-n', type=str, required=True, help='input pattern name')
    parser.add_argument('--dump_input', '-d', type=int, default=0, help='only dump input data. if 1, read shape from model; if 2, read shape from args (when simulator)')
    parser.add_argument('--input_shape', '-s', type=str, required=False, help='shape of input, only use when dump_input=2')
    parser.add_argument('--input_tensor_name', '-in', type=str, required=False, help='input tensor name, only use when dump_input=2')
    parser.add_argument('--input_pic_dir', '-id', type=str, required=False, help='input pic dir, only use when dump_input=2')
    parser.add_argument('--input_pic_shape', type=str, required=False, help='input pic shape, only use when dump_input=2')
    parser.add_argument('--input_bin_type', type=int, required=False, help='input bin type, only use when dump_input=2 and src/bintype=6')
    parser.add_argument('--input_frac_bit', type=int, required=False, help='input frac bit, only use when dump_input=2 and src/bintype=6')
    parser.add_argument('--gen_config', type=str, required=False, help='gen config dir')
    parser.add_argument('--sim_config', type=str, required=False, help='sim config dir')

    return parser.parse_args()


def rmdir(dir):
    if os.path.isdir(dir):
        for file in os.listdir(dir):
            rmdir(os.path.join(dir, file))
        if os.path.exists(dir):
            os.rmdir(dir)
    else:
        if os.path.exists(dir):
            os.remove(dir)


def read_config(file_path):
    config_file = open(file_path, "r", encoding='UTF-8-sig')
    config_data = config_file.readlines()

    configs = {}
    config_list = []
    for line in config_data:
        if line.isspace():
            continue
        if line != '\n' and line[0] != '#':
            config = line.split('=')
            for i in range(len(config)):
                config[i] = config[i].replace(' ', '')
                config[i] = config[i].replace('[', '')
                config[i] = config[i].replace(']', '')
                config[i] = config[i].replace('\n', '')
                config[i] = config[i].replace('\\', '/')
            if config[0] == 'blob_name':
                config_list.append(configs)
                configs = {}
            configs[config[0]] = config[1]
    config_list.append(configs)
    return config_list


def check_is_novaonnx(onnx_model):
    if onnx_model.producer_name == 'Novatek NovaOnnx Converter':
        return True
    else:
        attr = onnx_model.graph.node[0].attribute
        for j in range(len(attr)):
            if attr[j].name == 'layer_idx':
                return True
    return False


def read_mean_from_binaryproto(file_path):
    mean_blob = caffe_pb2.BlobProto()
    mean_blob.ParseFromString(open(file_path, 'rb').read())
    data = np.array(mean_blob.data)
    data = data.reshape(mean_blob.num, mean_blob.channels, mean_blob.height, mean_blob.width)
    return data


def read_img_file(ref_list_path, config_dir, ref_img_dir):
    image_info = open(ref_list_path).readlines()
    image_list = list(filter(None, image_info[0].split(';')))
    # img_file = config_dir + "/" + args['path/ref_img_dir'] + '/' + image_list[0].replace('\n', '')
    img_file_list = []
    for image_name in image_list:
        img_file = config_dir + "/" + ref_img_dir + '/' + image_name.replace('\n', '').replace(' ', '')
        img_file_list.append(img_file)

    return img_file_list


def multiplyList(List) :
    result = 1
    for x in List:
         result = result * x  
    return result 


def from_NV12(yuv_data, pic_shape):
    img_height = int(pic_shape[0])
    img_width = int(pic_shape[1])
    y_size = int(img_width * img_height)

    uv_height = int(img_height / 2)
    uv_width = int(img_width / 2)
    uv_size = int(uv_width * uv_height)

    Y = np.zeros((img_height, img_width), dtype=np.uint8)
    U = np.zeros((uv_height, uv_width), dtype=np.uint8)
    V = np.zeros((uv_height, uv_width), dtype=np.uint8)

    y_start = 0
    u_v_start = y_start + y_size
    u_v_end = u_v_start + (uv_size * 2)
    u_v_size = int(uv_width * uv_height)

    Y[:, :] = yuv_data[y_start : u_v_start].reshape((img_height, img_width))
    U_V = yuv_data[u_v_start : u_v_end].reshape((u_v_size, 2))

    U[:, :] = U_V[:, 0].reshape((uv_height, uv_width))
    V[:, :] = U_V[:, 1].reshape((uv_height, uv_width))

    return Y, U, V


def img2gray(images):
    images = images.astype(float).transpose(1,2,0)
    return np.dot(images, [[0.2989],[0.5870],[0.1140]]).transpose(2,0,1)


def bgr2YUV420(pic):
    # pic = yuv_pad(pic)
    pic = cv2.cvtColor(pic, cv2.COLOR_BGR2YUV)

    Y = pic[:,:,0]
    U = pic[:,:,1]
    V = pic[:,:,2]

    u_channel = []
    v_channel = []

    for w in range(U.shape[0]):
        for h in range(U.shape[1]):
            if w % 2 == 0 and h % 2 == 0:
                u_val = (int(U[w, h]) + int(U[w+1, h]) + int(U[w, h+1]) + int(U[w+1, h+1]) + 2) / 4
                u_val_int = int(u_val)

                v_val = (int(V[w, h]) + int(V[w+1, h]) + int(V[w, h+1]) + int(V[w+1, h+1]) + 2) / 4
                v_val_int = int(v_val)

                u_channel.append(u_val_int)
                v_channel.append(v_val_int)

    u_channel = np.array(u_channel)
    v_channel = np.array(v_channel)

    u_channel = u_channel.reshape([int(pic.shape[0]/2), int(pic.shape[1]/2)])
    v_channel = v_channel.reshape([int(pic.shape[0]/2), int(pic.shape[1]/2)])

    return Y, u_channel, v_channel


def yuv2bgr(Y, U, V, pic_shape):
    rgb_data = np.zeros((pic_shape[0], pic_shape[1], 3), dtype=np.uint8)
    V = np.repeat(V, 2, 0)
    V = np.repeat(V, 2, 1).astype(np.float32)
    U = np.repeat(U, 2, 0)
    U = np.repeat(U, 2, 1).astype(np.float32)
    Y = Y.astype(np.float32)

    # formula 1
    # r = Y + 0.00293*(U-128) + 1.39966*(V-128)
    # g = Y - 0.34399*(U-128) - 0.71729*(V-128)
    # b = Y + 1.77222*(U-128) - 0.00635*(V-128)

    # int formula 1
    r = (( 4096*Y + 12  *(U-128) + 5733*(V-128) ) /4096) + 0.5
    g = (( 4096*Y - 1409*(U-128) - 2938*(V-128) ) /4096) + 0.5
    b = (( 4096*Y + 7259*(U-128) - 26  *(V-128) ) /4096) + 0.5

    # formula 2
    # r = 1.164*(Y-16) + 1.596*(V-128)
    # g = 1.164*(Y-16) - 0.391*(U-128) - 0.813*(V-128)
    # b = 1.164*(Y-16) + 2.018*(U-128)

    # formula 3
    # r = 1.164*(Y-16) + 1.792*(V-128)
    # g = 1.164*(Y-16) - 0.213*(U-128) - 0.534*(V-128)
    # b = 1.164*(Y-16) + 2.114*(U-128)

    # formula 4
    # r = Y + (1.4022*(V-128))
    # g = Y - (0.3456*(U-128)) - (0.7145*(V-128))
    # b = Y + (1.7710*(U-128))

    r = np.where(r < 0, 0, r)
    r = np.where(r > 255, 255, r)

    g = np.where(g < 0, 0, g)
    g = np.where(g > 255, 255, g)

    b = np.where(b < 0, 0, b)
    b = np.where(b > 255, 255, b)

    rgb_data[:, :, 2] = b
    rgb_data[:, :, 1] = g
    rgb_data[:, :, 0] = r

    bgr_data = rgb_data[:, :, [2, 1, 0]]
    # cv2.imshow("bgr incorrect colors", bgr_data)
    # cv2.waitKey(0)

    return bgr_data


def read_bin_pic(in_type, img_file):  # read bin data when module test (h*w*c)
    if in_type == '0':
        feature = np.fromfile(img_file, dtype=np.int8)
    elif in_type == '1':
        feature = np.fromfile(img_file, dtype=np.uint8)
    elif in_type == '2':
        feature = np.fromfile(img_file, dtype=np.int16)
    elif in_type == '3':
        feature = np.fromfile(img_file, dtype=np.uint16)
    elif in_type == '6':
        feature = np.fromfile(img_file, dtype=np.float32)
    else:
        assert False, "Unsupported input type"
    
    return feature


def read_bin_data(sim_configs, sim_input_dir, input_bin_type=None):
    if sim_configs['src/imgfmt'] == '0':  # jpg img fmt ,only uint8
        sim_configs['src/bintype'] = '1'

    if sim_configs['src/bintype'] == '0':
        bin_data = np.fromfile(sim_input_dir, dtype=np.int8)
    elif sim_configs['src/bintype'] == '1':
        bin_data = np.fromfile(sim_input_dir, dtype=np.uint8)
    elif sim_configs['src/bintype'] == '2':
        bin_data = np.fromfile(sim_input_dir, dtype=np.int16)
    elif sim_configs['src/bintype'] == '3':
        bin_data = np.fromfile(sim_input_dir, dtype=np.uint16)
    elif sim_configs['src/bintype'] == '6':
        if input_bin_type == 0:
            bin_data = np.fromfile(sim_input_dir, dtype=np.int8)
        elif input_bin_type == 1:
            bin_data = np.fromfile(sim_input_dir, dtype=np.uint8)
        elif input_bin_type == 2:
            bin_data = np.fromfile(sim_input_dir, dtype=np.int16)
        elif input_bin_type == 3:
            bin_data = np.fromfile(sim_input_dir, dtype=np.uint16)
        elif input_bin_type == 4:
            bin_data = np.fromfile(sim_input_dir, dtype=np.uint32)
        elif input_bin_type == 5:
            bin_data = np.fromfile(sim_input_dir, dtype=np.int32)

    else:
        assert False, "Unsupported bintype!"
    
    return bin_data


def read_bgr_from_yuv420(sim_configs, sim_input_dir, simin_shape, index):
    pic = read_bin_data(sim_configs, sim_input_dir)

    # multi-batch
    if simin_shape[-1] != 1:
        pic_shape = [int(simin_shape[1]), int(simin_shape[0])]
        pic = np.reshape(pic, [simin_shape[-1], -1])
        pic = pic[index, :]
        Y, U, V = from_NV12(pic, pic_shape)
        bgr_data = yuv2bgr(Y, U, V, pic_shape)
    else:    
        pic_shape = [int(simin_shape[1]), int(simin_shape[0])]
        Y, U, V = from_NV12(pic, pic_shape)
        bgr_data = yuv2bgr(Y, U, V, pic_shape)

    return bgr_data


def read_yonly_from_bin(sim_configs, sim_input_dir, simin_shape, index):
    pic = read_bin_data(sim_configs, sim_input_dir)

    pic_shape = [int(simin_shape[1]), int(simin_shape[0])]

    # multi-batch
    if simin_shape[-1] != 1:
        pic = np.reshape(pic, [int(simin_shape[3]), -1])[index, :]  # only get index batch
        
    if pic.size > simin_shape[2] * simin_shape[1] * simin_shape[0]:  # only read clip
        pic = pic[0: simin_shape[2] * simin_shape[1] * simin_shape[0]]

    yonly_data = np.reshape(pic, pic_shape)
    return yonly_data


def read_pic_from_bin(sim_configs, sim_input_dir, simin_shape, index):
    pic = read_bin_data(sim_configs, sim_input_dir)
    pic_shape = [int(simin_shape[2]), int(simin_shape[1]), int(simin_shape[0])]
    # multi-batch
    if simin_shape[-1] != 1:
        pic = np.reshape(pic, [int(simin_shape[3]), -1])[index, :]  # only get index batch

    if pic.size > simin_shape[2] * simin_shape[1] * simin_shape[0]:  # only read clip
        pic = pic[0: simin_shape[2] * simin_shape[1] * simin_shape[0]]
        
    pic_data = np.reshape(pic, pic_shape).transpose(1,2,0)
    return pic_data


def read_feature_from_bin(sim_configs, sim_input_dir, input_bin_type):
    return read_bin_data(sim_configs, sim_input_dir, input_bin_type)


def yuv_pad(feature):
    if (feature.shape[0]%2) != 0:
        feature = np.pad(feature, ((0,1), (0,0), (0,0)), 'constant')
    if (feature.shape[1]%2) != 0:
        feature = np.pad(feature, ((0,0), (0,1), (0,0)), 'constant')
    return feature


def yuv_crop(feature):
    if (feature.shape[0]%2) != 0:
        feature = feature[:-1, :, :]
    if (feature.shape[1]%2) != 0:
        feature = feature[:, :-1, :]
    return feature


def del_layer_idx(graph):
    for i in range(len(graph.node)):
        attr = graph.node[i].attribute
        for j in range(len(attr)):
            if attr[j].name == 'layer_idx':
                attr.pop(j)
                break
        if graph.node[i].op_type == 'AveragePool' or graph.node[i].op_type == 'MaxPool':
            for j in range(len(attr)):
                if attr[j].name == 'pool_at_pad':
                    attr.pop(j)
                    break 
    return graph


def modify_name(graph):
    for input in graph.input:
        input.name = input.name.replace('/','_')
        input.name = input.name.replace(':','_')

    init_name_list = []
    for initializer in graph.initializer:
        initializer.name = initializer.name.replace('/','_')
        initializer.name = initializer.name.replace(':','_')
        init_name_list.append(initializer.name)

    for output in graph.output:
        output.name = output.name.replace('/', '_')
        output.name = output.name.replace(':', '_')

    for value_info in graph.value_info:
        value_info.name = value_info.name.replace('/', '_')
        value_info.name = value_info.name.replace(':', '_')

    for node in graph.node:
        node.name = node.name.replace('/', '_')
        node.name = node.name.replace(':', '_')
        for i in range(len(node.input)):
            node.input[i] = node.input[i].replace('/', '_')
            node.input[i] = node.input[i].replace(':', '_')
        for i in range(len(node.output)):
            node.output[i] = node.output[i].replace('/', '_')
            node.output[i] = node.output[i].replace(':', '_')

    name_dict = {}
    # modify Conv weight name
    for i in range(len(graph.node)):
        if graph.node[i].op_type == 'Conv':
            if graph.node[i].input[1] in init_name_list:
                name_dict.setdefault(graph.node[i].input[1],
                                     graph.node[i].op_type + "_" + graph.node[i].input[1] + "_W")
                graph.node[i].input[1] = graph.node[i].op_type + "_" + graph.node[i].input[1] + "_W"
            if len(graph.node[i].input) > 2:
                if graph.node[i].input[2] in init_name_list:
                    name_dict.setdefault(graph.node[i].input[2],
                                         graph.node[i].op_type + "_" + graph.node[i].input[2] + "_B")
                    graph.node[i].input[2] = graph.node[i].op_type + "_" + graph.node[i].input[2] + "_B"

        # modify output tensor_name to (node_name)_Y
        for k in range(len(graph.node[i].input)):
            if graph.node[i].input[k] in name_dict:
                graph.node[i].input[k] = name_dict[graph.node[i].input[k]]
        for l in range(len(graph.node[i].output)):
            name_dict.setdefault(graph.node[i].output[l], graph.node[i].op_type + "_" + graph.node[i].output[l] + "_Y")
            graph.node[i].output[l] = graph.node[i].op_type + "_" + graph.node[i].output[l] + "_Y"

    # modify graph output tensor_name to (node_name)_Y
    for m in range(len(graph.output)):
        if graph.output[m].name in name_dict:
            graph.output[m].name = name_dict[graph.output[m].name]

    # modify value info name
    for n in range(len(graph.value_info)):
        if graph.value_info[n].name in name_dict:
            graph.value_info[n].name = name_dict[graph.value_info[n].name]

    # modify input name
    for o in range(len(graph.input)):
        if graph.input[o].name in name_dict:
            graph.input[o].name = name_dict[graph.input[o].name]

            # modify initializer name
    for p in range(len(graph.initializer)):
        if graph.initializer[p].name in name_dict:
            graph.initializer[p].name = name_dict[graph.initializer[p].name]

    return graph


def remove_initializer_from_input(model):
    # if model.ir_version < 4:
    #     print("Model with ir_version below 4 requires to include initilizer in graph input")
    #     return

    inputs = model.graph.input
    name_to_input = {}
    for input in inputs:
        name_to_input[input.name] = input

    for initializer in model.graph.initializer:
        if initializer.name in name_to_input:
            inputs.remove(name_to_input[initializer.name])

    return model


def bilinear_interpolate(src, shape): 
    # if shape[0] != 1:  # only support batch=1
    #     shape[0] = 1
    
    if len(src.shape)==2:  # reshape to h*w*c
        src = np.reshape(src, [src.shape[0], src.shape[1], 1]) 

    dst_size =  [shape[2], shape[3]]

    height_src, width_src, channel_src = src.shape  # (h, w, ch)
    height_dst, width_dst = dst_size  # (h, w)

    # Align the center, project the coordinate of the target image onto the original image
    # ws_p = np.array([(i + 0.5) / width_dst * width_src - 0.5 for i in range(width_dst)], dtype=np.float32)
    # hs_p = np.array([(i + 0.5) / height_dst * height_src - 0.5 for i in range(height_dst)], dtype=np.float32)
    ws_p = np.array([i * ((width_src-1)/(width_dst-1)) for i in range(width_dst)], dtype=np.float32)
    hs_p = np.array([i * ((height_src-1)/(height_dst-1)) for i in range(height_dst)], dtype=np.float32)
    ws_p = np.repeat(ws_p.reshape(1, width_dst), height_dst, axis=0)
    hs_p = np.repeat(hs_p.reshape(height_dst, 1), width_dst, axis=1)

    # Find the coordinates of the nearest neighbor points
    ws_0 = np.clip(np.floor(ws_p), 0, width_src - 2).astype(np.int32)
    hs_0 = np.clip(np.floor(hs_p), 0, height_src - 2).astype(np.int32)
    ws_1 = ws_0 + 1
    hs_1 = hs_0 + 1

    # Pixel value of the four nearest neighbor points
    f_00 = src[hs_0, ws_0, :].T
    f_01 = src[hs_0, ws_1, :].T
    f_10 = src[hs_1, ws_0, :].T
    f_11 = src[hs_1, ws_1, :].T

    # cal weight
    w_00 = ((hs_1 - hs_p) * (ws_1 - ws_p)).T
    w_01 = ((hs_1 - hs_p) * (ws_p - ws_0)).T
    w_10 = ((hs_p - hs_0) * (ws_1 - ws_p)).T
    w_11 = ((hs_p - hs_0) * (ws_p - ws_0)).T

    # cal target image pixel value
    return (f_00 * w_00).T + (f_01 * w_01).T + (f_10 * w_10).T + (f_11 * w_11).T


def inputs_preproc(feature, args, config_dir, sim_configs=None):
    if 'preproc/meansub/en' in args:
        if args['preproc/meansub/en'] != '0':
            if args['preproc/meansub/fmt'] == '0':  # TXT
                if args['preproc/meansub/mode'] == '0':  # dc mode
                    mean_file = open(config_dir + "/" + args['path/mean_path'], "r")
                    mean_data = mean_file.readline().split(',')
                    mean_data = list(map(float, mean_data))  # str2float

                if args['preproc/meansub/mode'] == '1':  # Plannar mode
                    fname = config_dir + "/" + args['path/mean_path']
                    with open(fname, 'r+') as f:
                        mean_data = [i[:-1].split(',') for i in f.readlines()]

                    for i in range(len(mean_data)):
                        if mean_data[i][-1] == ' ':
                            mean_data[i] = mean_data[i][:-1]
                        mean_data[i] = list(map(float, mean_data[i]))
                    if len(mean_data)==1:
                        mean_data = mean_data[0]
                        shape = list(map(int, [mean_data[2], mean_data[1], mean_data[0]]))  # get shape
                        mean_data = np.reshape(mean_data[3:], shape)
                    elif len(mean_data)==2:
                        shape = list(map(int, [mean_data[0][2], mean_data[0][1], mean_data[0][0]]))
                        mean_data = np.reshape(mean_data[1], shape)

                    mean_data = bilinear_interpolate(mean_data.transpose(1,2,0), feature.shape)
                    # mean_data = cv2.resize(mean_data.transpose(1,2,0), (feature.shape[3], feature.shape[2]))
                    if len(mean_data.shape)==2:
                        mean_data = np.expand_dims(mean_data, axis=-1)
                    mean_data = mean_data.transpose(2, 0, 1).astype(np.float32)

            elif args['preproc/meansub/fmt'] == '1':  # BINARYPROTO
                mean_data = read_mean_from_binaryproto(config_dir + "/" + args['path/mean_path'])
                mean_data = mean_data.astype(np.float32)
                if args['preproc/meansub/mode'] == '1':  # Plannar
                    channel_num = mean_data.shape[1]
                    mean_data = mean_data[0,:,:,:]
                    # mean_data = cv2.resize(mean_data.transpose(1, 2, 0), (feature.shape[3], feature.shape[2]))
                    mean_data = bilinear_interpolate(mean_data.transpose(1,2,0), feature.shape)
                    mean_data = np.reshape(mean_data.transpose(2, 0, 1), (channel_num, feature.shape[2], feature.shape[3]))

            else:
                assert False, "Unsupported meansub_fmt!"

            if args['preproc/meansub/fmt_type'] == '0':
                if args['preproc/meansub/mode'] == '1':  # FMT_YONLY and Plannar mode
                    if mean_data.shape[0] != 1:
                        if args['preproc/meansub/fmt_type'] == '4':
                            mean_data = img2gray(mean_data)
                        elif args['preproc/meansub/fmt_type'] == '1':
                            mean_data = [mean_data[2], mean_data[1], mean_data[0]]  # rgb2bgr
                            mean_data = img2gray(mean_data)
                        else:
                            mean_data = mean_data[0:feature.shape[1], :, :]
                    if len(mean_data.shape) == 4:
                        mean_data = mean_data[0,:,:,:]
                    mean_data = np.reshape(mean_data, feature.shape)

            elif args['preproc/meansub/fmt_type'] == '1':  # FMT_RGB
                if len(mean_data) == 1:
                    mean_data = [mean_data[0], mean_data[0], mean_data[0]]
                if args['preproc/out_fmt'] == '3' and args['preproc/meansub/fmt'] == '0':  # RGB2BGR
                    mean_data = [mean_data[2], mean_data[1], mean_data[0]]
                if args['preproc/out_fmt'] == '3' and args['preproc/meansub/fmt'] == '1':  # RGB2BGR and binaryproto
                    mean_data = mean_data[[2, 1, 0], :, :]

            elif args['preproc/meansub/fmt_type'] == '3':  # FMT_FEAT
                pass

            elif args['preproc/meansub/fmt_type'] == '4':  # FMT_BGR
                if len(mean_data) == 1:
                    mean_data = [mean_data[0], mean_data[0], mean_data[0]]
                if args['preproc/out_fmt'] == '2' and args['preproc/meansub/fmt'] == '0':  # BGR2RGB
                    mean_data = [mean_data[2], mean_data[1], mean_data[0]]
                if args['preproc/out_fmt'] == '2' and args['preproc/meansub/fmt'] == '1':  # BGR2RGB and binaryproto
                    mean_data = mean_data[[2, 1, 0], :, :]

            else:
                assert False, "Unsupported meansub_fmt_type"

            if args['preproc/meansub/mode'] == '0':  # meansub DC mode
                for i in range(feature.shape[1]):
                    feature[:, i, :, :] -= mean_data[i]
            elif args['preproc/meansub/mode'] == '1':  # meansub Plannar mode
                mean_data = np.array(mean_data)
                feature = feature - mean_data
            else:
                assert False, "Unsupported meansub_mode"

            # do clamp
            if not (args['preproc/in/fmt'] == '3' and args['preproc/out_fmt'] == '4') and sim_configs==None:
                feature = np.where(feature < -128, -128, feature)
                feature = np.where(feature > 127, 127, feature)

    if 'preproc/meansub/en' in args:
        if args['preproc/normalize/en'] == '1':
            norm_scale = list(map(float, args['preproc/normalize/scale'].split(',')))
            if len(norm_scale) == 1:
                feature *= float(args['preproc/normalize/scale'])
            # elif sim_configs == None:  # per channel norm (skip when sim)
            elif len(norm_scale) == feature.shape[1]:
                for i in range(len(norm_scale)):
                    feature[:,i,:,:] *= norm_scale[i]

    return feature


def load_data(input_shape, configs, config_dir, sim_configs=None, sim_input_dir=None, output_dir=None, simin_shape=None, input_bin_type=None, input_frac_bit=None):
    img_file_list = read_img_file(config_dir + "/" + configs['path/ref_list_path'], config_dir, configs['path/ref_img_dir'])  
    feature_out = np.ones(input_shape).astype(np.float32)

    if input_shape[0] != 1 and len(img_file_list) == 1:  # repeat input image
        img_file_list = img_file_list * input_shape[0]

    for i in range(input_shape[0]):  # multi batch
        img_file = img_file_list[i]
        if configs['preproc/in/fmt'] == '0' or configs['preproc/out_fmt'] == '0':
            if sim_configs:  # read from sdk input
                pic = read_yonly_from_bin(sim_configs, sim_input_dir, simin_shape, i)  
            else:  # read from gen_config
                pic = cv2.imread(img_file, cv2.IMREAD_UNCHANGED)
                if len(pic.shape) == 3:
                    if pic.shape[2] == 3:
                        pic = cv2.cvtColor(pic, cv2.COLOR_BGR2GRAY)
                    elif pic.shape[2] == 4:
                        pic = cv2.cvtColor(pic, cv2.COLOR_BGRA2GRAY)
                else:
                    pic = pic
            pic = yuv_pad(pic.reshape(pic.shape[0], pic.shape[1], 1))
            dst = bilinear_interpolate(pic, input_shape)
            feature = np.reshape(dst.transpose(2, 0, 1), [1, input_shape[1], input_shape[2], input_shape[3]])
            feature = inputs_preproc(feature.astype(np.float32), configs, config_dir, sim_configs)

        elif configs['preproc/in/fmt'] == '1':  # FMT_RGB
            if sim_configs:  # read from sdk input
                pic = read_pic_from_bin(sim_configs, sim_input_dir, simin_shape, i)
                if configs['preproc/out_fmt'] == '3':  # BGR out
                    pic = pic[:, :, [2, 1, 0]]  # BGR TO RGB
            else:  # read from gen_config
                pic = cv2.imread(img_file)  # cv2.read get BGR
                if pic is None:
                    pic = read_bin_pic(configs['preproc/in/type'], img_file)
                    pic = pic.reshape([input_shape[1], input_shape[2], input_shape[3]]).transpose(1, 2, 0)
                if configs['preproc/out_fmt'] == '2':  # RGB out
                    pic = pic[:, :, [2, 1, 0]]  # BGR TO RGB
            pic = yuv_pad(pic)
            dst = bilinear_interpolate(pic, input_shape)
            feature = np.reshape(dst.transpose(2, 0, 1), [1, input_shape[1], input_shape[2], input_shape[3]])
            feature = inputs_preproc(feature.astype(np.float32), configs, config_dir, sim_configs)

        elif configs['preproc/in/fmt'] == '4':  # FMT_BGR
            if sim_configs:  # read from sdk input
                pic = read_pic_from_bin(sim_configs, sim_input_dir, simin_shape, i)
            else:  # read from gen_config
                pic = cv2.imread(img_file)
            pic = yuv_pad(pic)
            if configs['preproc/out_fmt'] == '2':  # RGB out
                pic = pic[:, :, [2, 1, 0]]  # BGR TO RGB
            dst = bilinear_interpolate(pic, input_shape)
            feature = np.reshape(dst.transpose(2, 0, 1), [1, input_shape[1], input_shape[2], input_shape[3]])
            feature = inputs_preproc(feature.astype(np.float32), configs, config_dir, sim_configs)

        elif configs['preproc/in/fmt'] == '2':  # FMT_YUV420
            if sim_configs:
                pic = read_bgr_from_yuv420(sim_configs, sim_input_dir, simin_shape, i)
            else:
                pic = cv2.imread(img_file)
                pic = yuv_pad(pic)
                Y, U, V = bgr2YUV420(pic)
                pic = yuv2bgr(Y, U, V, [int(pic.shape[0]), int(pic.shape[1])])
                # pic.tofile('pic' + ".bin")

            if configs['preproc/out_fmt'] == '2' or configs['preproc/out_fmt'] == '3':
                pic = yuv_pad(pic)

            if configs['preproc/out_fmt'] == '2':  # RGB out
                pic = pic[:, :, [2, 1, 0]]  # BGR TO RGB

            dst = (bilinear_interpolate(pic, input_shape) + 0.5).astype(int)
            feature = np.reshape(dst.transpose(2, 0, 1), [1, input_shape[1], input_shape[2], input_shape[3]])
            feature = inputs_preproc(feature.astype(np.float32), configs, config_dir, sim_configs)

        elif configs['preproc/in/fmt'] == '3' or configs['preproc/out_fmt'] == '4':  # feature mode
            if sim_configs:
                feature = read_feature_from_bin(sim_configs, sim_input_dir, input_bin_type)
                if simin_shape[-1] != 1:
                    # get i-th batch data
                    feature = np.reshape(feature, [simin_shape[-1], -1])[i, :]
            else:
                feature = read_bin_pic(configs['preproc/in/type'], img_file)
                if input_shape[0] != 1 and multiplyList(input_shape) == feature.size:
                    # if multi-batch and bin size = input batch*c*h*w
                    # get i-th batch data
                    feature = np.reshape(feature, [input_shape[0], -1])[i, :]  
            if sim_configs:
                if sim_configs["src/bintype"] != '6':
                    input_frac_bit = int(configs["preproc/in/frac_bit_num"])
                    
            if configs['preproc/in/type'] != '6' or sim_configs:  # do frac bit when sim preproc or not float in.
                if input_frac_bit:
                    feature = feature / (2 ** int(input_frac_bit))
                else: 
                    feature = feature / (2 ** int(configs["preproc/in/frac_bit_num"]))
            feature = feature.astype(np.float32)

            if input_shape[0] != 1 and len(img_file_list) == 1:  # multi batch form 1 bin
                feature = np.reshape(feature, [input_shape[0], -1])
                for j in range(input_shape[0]):
                    feature_batch = feature[j, :]
                    feature_batch = inputs_preproc(feature_batch.reshape([1, input_shape[1], input_shape[2], input_shape[3]]), configs, config_dir, sim_configs)
                    feature_out[j,:,:,:] = feature_batch
                return feature_out

            if len(input_shape) == 4:
                if feature.size > input_shape[1]*input_shape[2]*input_shape[3]:
                    feature = feature[0: input_shape[1]*input_shape[2]*input_shape[3]]
                feature = inputs_preproc(feature.reshape([1, input_shape[1], input_shape[2], input_shape[3]]), configs, config_dir, sim_configs)
            elif len(input_shape) == 3:  # sinonnx module test
                if feature.size > input_shape[0]*input_shape[1]*input_shape[2]:
                    feature = feature[0: input_shape[0]*input_shape[1]*input_shape[2]]
                feature = inputs_preproc(feature.reshape([1, input_shape[0], input_shape[1], input_shape[2]]), configs, config_dir, sim_configs)
                return feature[0, :, :, :]
            elif len(input_shape) == 2:  # sinonnx module test
                if (input_shape[0] * input_shape[1]) == (feature.size):
                    feature = inputs_preproc(feature.reshape([1, 1, input_shape[0], input_shape[1]]), configs, config_dir, sim_configs)
                    feature_out[:,:] = feature[0, 0, :, :]
                    return feature_out

                elif input_shape[0] == len(img_file_list):
                    feature = inputs_preproc(feature.reshape([1, 1, 1, input_shape[1]]), configs, config_dir, sim_configs)
                    feature_out[i,:] = feature[0, 0, 0, :]
                    continue
                else:
                    feature = inputs_preproc(feature.reshape([1, 1, input_shape[0], input_shape[1]]), configs, config_dir, sim_configs)
                    return feature[0, 0, :, :]
            elif len(input_shape) == 1:  # sinonnx module test
                feature = inputs_preproc(feature.reshape([1, 1, 1, input_shape[0]]), configs, config_dir, sim_configs)
                return feature[0, 0, 0, :]
        else:
            assert False, "Unsupported in_fmt"

        feature_out[i,:,:,:] = feature
    return feature_out


def get_tensor_shape(tensor):
    dims = tensor.type.tensor_type.shape.dim
    n = len(dims)
    return [dims[i].dim_value for i in range(n)]


def getInputShape(input_config, input_shape):
    if 'dynamic_batch/en' in input_config:
        if input_config['dynamic_batch/en'] == '1':
            batchs = input_config['dynamic_batch/batch_size'].split(',')
            batchs = list(map(int, batchs))  # str2ints

            input_shape[0] = max(batchs)
    
    if 'multiscale/en' in input_config:
        if input_config['multiscale/en'] == '1':
            widths = input_config['multiscale/width'].split(',')
            widths = list(map(int, widths))  # str2ints
            input_shape[3] = max(widths)

            height = input_config['multiscale/height'].split(',')
            height = list(map(int, height))  # str2ints
            input_shape[2] = max(height)

    return input_shape


# dump preproc result, shape from args
def dump_preproc_result(pattern_name, config_dir, configs, input_shape, input_name, sim_configs=None, input_pic_dir=None, input_pic_shape=None, input_bin_type=None, input_frac_bit=None):
    output_dir = config_dir + "/" + configs[0]['path/out_dir'] + "/" + pattern_name + "/debug/golden_tensor/"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    input_shape = list(filter(None, input_shape.split(',')))
    input_shape = list(map(int,input_shape))
    input_pic_shape = list(filter(None, input_pic_shape.split(',')))
    input_pic_shape = list(map(int,input_pic_shape))

    inputs = {}
    if len(configs) == 1 or len(configs) == 2:  # only one input (gen config)
        if len(configs) == 1:
            input_config = configs[0]
        else:  # if set blob_name
            input_config = configs[1]

        if len(sim_configs) == 1:
            input_sim_config = sim_configs[0]
        else:
            input_sim_config = sim_configs[1]
        
        # if set multiscale or dynamic_batch
        input_shape = getInputShape(input_config, input_shape)

        image = load_data(input_shape, input_config, config_dir, input_sim_config, input_pic_dir, output_dir, input_pic_shape, input_bin_type, input_frac_bit)
        inputs[input_name] = image
    else:  # Multiple inputs
        for i, config in enumerate(configs[1:]):
            if input_name == config['blob_name']:
                input_config = configs[i+1]
                input_sim_config = sim_configs[i+1]
                image = load_data(input_shape, input_config, config_dir, input_sim_config, input_pic_dir, output_dir, input_pic_shape, input_bin_type, input_frac_bit)
                inputs[config['blob_name']] = image
    # for i in range(input_num):
    name = input_name.replace('/', '_')
    name = name.replace(':', '_')
    inputs[input_name].tofile(output_dir + name + ".bin")


def runtime_infer(input_dir, pattern_name, config_dir, configs, dump_input):
    output_dir = config_dir + "/" + configs[0]['path/out_dir'] + "/" + pattern_name + "/debug/ref_layer/fronted/"
    rmdir(output_dir)
    os.makedirs(output_dir)

    onnx_model = onnx.load(input_dir)
    graph = onnx_model.graph

    is_novaonnx = check_is_novaonnx(onnx_model)

    if is_novaonnx:
        onnx_model.opset_import[0].version = 18
        # print("This model is a nova onnx model")
        graph = del_layer_idx(graph)
    else:
        # print("This model is not a nova onnx model")
        graph = modify_name(graph)

    if onnx_model.ir_version >= 4 or dump_input != 0:
        # onnx runtime Model with ir_version below 4 requires to include initilizer in graph input
        # remove init from input if only dump_input
        remove_initializer_from_input(onnx_model)

    input_shape = []
    input_name = []
    input_num = len(graph.input)

    # add input to output, and save input shape
    for i in range(input_num):
        graph.output.insert(i, graph.input[i])
        shape = get_tensor_shape(graph.input[i])
        # print(shape)
        input_shape.append(shape)
        input_name.append(graph.input[i].name)

    # load data from data_file
    inputs = {}
    if len(configs) == 1:  # only one input
        image = load_data(input_shape[0], configs[0], config_dir, output_dir=output_dir)
        inputs[input_name[0]] = image
    else:  # Multiple inputs
        for i, config in enumerate(configs[1:]):
            for j, name in enumerate(input_name):
                if name == config['blob_name']:
                    shape = input_shape[j]
            image = load_data(shape, config, config_dir, output_dir=output_dir)
            inputs[config['blob_name']] = image

    # dump preproc result, shape from model
    if dump_input == 1:
        for i in range(input_num):
            name = input_name[i].replace('/', '_')
            name = name.replace(':', '_')
            inputs[input_name[i]].tofile(output_dir + name + ".bin")
        return

    # save all node output
    for node in graph.node:
        for output in node.output:
            graph.output.extend([onnx.ValueInfoProto(name=output)])

    model_file = output_dir + "temp.onnx"
    onnx.save(onnx_model, model_file)

    # do inference
    sess = rt.InferenceSession(model_file)
    outputs = {}

    for out in sess.get_outputs():
        tensor = sess.run([out.name], inputs)
        onnx_out = tensor[0]
        name = out.name.replace('/', '_')
        name = name.replace(':', '_')
        onnx_out.tofile(output_dir + name + ".bin")
        outputs[str(out.name)] = tensor

    # save weight
    weights, names = [], []
    for init in graph.initializer:
        weights.append(numpy_helper.to_array(init))
        names.append(init.name)

    for i in range(len(weights)):
        name = names[i].replace('/', '_')
        name = name.replace(':', '_')
        weights[i].tofile(output_dir + name + ".bin")

    os.remove(model_file)
    return outputs


if __name__ == "__main__":
    args = process_command()

    if args.gen_config:
        gen_config_dir = args.gen_config
        sim_config_dir = args.sim_config
    else:  # when module test
        gen_config_dir = args.config_dir + "/../nvtai_tool/config/" + args.name + "/cnn25/gen_config.txt"
        sim_config_dir = args.config_dir + "/../nvtai_tool/config/" + args.name + "/cnn25/sim_config.txt"
    configs = read_config(gen_config_dir)
    sim_configs = read_config(sim_config_dir)

    if args.dump_input == 2:
        # compare preproc cos & diff, when simulator stage
        dump_preproc_result(args.name, args.config_dir, configs, args.input_shape, args.input_tensor_name, sim_configs, args.input_pic_dir, args.input_pic_shape, args.input_bin_type, args.input_frac_bit)
    else:
        # run onnx model or get onnx model input shape, need import opencv & onnx
        import cv2
        import onnxruntime as rt
        import onnx
        from onnx import numpy_helper
        from onnx import helper
        from onnx import TensorProto

        runtime_infer(args.input, args.name, args.config_dir, configs, args.dump_input)
